import math
import random
from Env.action import *
from mylog import red_debug_print as debug_print

Uavid = list(range(13,19))
Uav_pos = [[4000, 7500], [3500, 8500], [3000, 6000], [3000, 7000], [1000, 7500], [2000, 9000]]
# 四辆坦克的序号
TankId = [1, 2, 3, 4]
Tanks_roads = [
    [[3920, 5960], [4020, 6600], [3900, 7450], [3000, 9100]],
    [[3825, 5933], [2225, 8340]],
    [[3753, 5890], [2193, 7978]],
    [[3634, 5851], [2070, 7894]],
]
TankRadomRoads= [
    [[2918, 8917], [3010, 8940], [3084, 9019], [2945, 9148], [2870, 9123]],
    [[2225, 8340], [2185, 8300], [2270, 8295], [2160, 8400], [2280, 8390]],
    [[2193, 7978], [2140, 7930], [2260, 7930], [2100, 8031], [2234, 8005]],
    [[2070, 7894], [2030, 7840], [2150, 7840], [2005, 7950], [2103, 7960]],
]

Uacid = [7,8,9,10,11,12]
Uac_roads = [[[3270+i*5, 5900+i*5], [2650+i*5,6500+i*5], [1950+i*5, 7000+i*5], [1533,8292], [1654+i*5, 8368+i*5]] for i in range(6)]

IcvId = [5,6]
Icv_roads = [[[3418+i*5, 6340+i*5], [3200+i*5,8000+i*5], [2250+i*5, 8750+i*5]] for i in range(2)]

Uac_pos = [[1660+i*5, 8410+i*5] for i in range(6)]
Icv_pos = [[2250+i*5, 8750+i*5] for i in range(2)]
IcvLastPos = [[2035 + i*5, 8685 + i*5] for i in range(2)]
Tank_pos = [[2910, 9066], [2225, 8340],  [2193, 7978],  [2070, 7894]]

UidMap = {
    'Uav': Uavid ,
    'Tan': TankId ,
    'Uac': Uacid,
    'Icv': IcvId,
}

class UnitStatus:
    def __init__(self, status, idx):
        self.status = status
        self.idx = idx

# 动作执行
class ActionGenerator:
    def __init__(self):
        print("Creat ActionGenerator Instance")
        self.state = None
        self.action_function_map = {
            "":                         self.NoneFunc,
            "UavDetection":             self.UavDetection,
            "UavUpping":                self.UavUp,
            "UavStop":                  self.UavStop,
            "TankGo":                   self.TankGo,
            "TankSupportWaiting":       self.NoneFunc,
            "TankRandomMoving":         self.TankRandomMove,
            "UacGo":                    self.UacGo,
            "IcvGo":                    self.IcvGo,
            "IcvWaiting":               self.NoneFunc,
            "IcvDown":                  self.IcvDown,
            "IcvLastMoving":            self.IcvLastMove,
            "AutoAttackAll":            self.AutoAttackAll,
        }

    def generate(self, status_list):
        self.cmds = []
        for status in status_list:
            self.action_function_map[status.status](status.idx)
        return [cmd for cmd in self.cmds if cmd] 

    def __calc_course(self, idx, poi):
        unit = self.state.unit_dict[idx]
        deg = math.atan2((poi[1]-unit.y), (poi[0]-unit.x)) / math.pi * 180
        return deg

    def isalive(self, uid):
        return self.state.unit_dict[uid].isalive

    def NoneFunc(self, idx):
        return
        
    def Distance(self, uid, poi):
        unit = self.state.unit_dict[uid]
        return ((unit.x - poi[0])**2 + (unit.y - poi[1])**2)**0.5
    #-------------------------无人机动作---------------------------
    def calc_course(self, idx, poi):
        unit = self.state.unit_dict[idx]
        deg = math.atan2((poi[1]-unit.y), (poi[0]-unit.x)) / math.pi * 180
        return deg
        
    # 无人机探测任务
    def UavDetection(self, idx):
        uavid = Uavid[idx]
        pos = Uav_pos[idx]
        if not self.isalive(uavid):
            return
        course = self.calc_course(uavid, pos)
        if course == 0:
            course = 1
        cmd = MoveDirecWrj(uavid, course)
        self.cmds.append(cmd)
    # 无人机升空
    def UavUp(self, idx):
        uavid = Uavid[idx]
        if not self.isalive(uavid):
            return
        cmd = HoverSelf(uavid, 1)
        self.cmds.append(cmd)
    # 无人机悬停
    def UavStop(self, idx):
        uavid = Uavid[idx]
        if not self.isalive(uavid):
            return
        cmd = HoverSelf(uavid, 0)
        self.cmds.append(cmd)

    def TankGo(self, idx):
        tank_id = TankId[idx]
        pois = Tanks_roads[idx]
        if not self.isalive(tank_id):
            return 
        cmd = MoveAsLineTank(tank_id, pois, 80)
        self.cmds.append(cmd)

    def TankRandomMove(self, idx):
        tank_id = TankId[idx]
        poi = random.choice(TankRadomRoads[idx])
        if not self.isalive(tank_id):
            return 
        cmd = MoveToPointTank(tank_id, poi[0], poi[1], 60)
        self.cmds.append(cmd)

    def UacGo(self, idx):
        uac_id = Uacid[idx]
        pois = Uac_roads[idx]
        if not self.isalive(uac_id):
            return
        cmd = MoveAsLineWrc(uac_id, pois, 80)
        self.cmds.append(cmd)
        
    def IcvGo(self, idx):
        icv_id = IcvId[idx]
        pois = Icv_roads[idx]
        if not self.isalive(icv_id):
            return
        cmd = MoveAsLineIcv(icv_id, pois, 100)
        self.cmds.append(cmd)
        
    def IcvLastMove(self, idx):
        icv_id = IcvId[idx]
        poi = IcvLastPos[idx]
        if not self.isalive(icv_id):
            return
        cmd = MoveToPointIcv(icv_id, poi[0], poi[1], 100)
        self.cmds.append(cmd)

    def AutoAttackAll(self, idx):
        import json
        # print("AutoAttackAll, Red Qbs is ")
        # print (json.dumps(self.state.state['red']['qb'], ensure_ascii=False))
        for unit in self.state.red.tanks:
            self.auto_attack(unit.uid)
        for unit in self.state.red.wrc:
            self.auto_attack(unit.uid)
        for unit in self.state.red.bbb:
            self.auto_attack(unit.uid)
    
    def calc_weight(self, unit, target):
        weight = -1
        dis = self.Distance(unit.uid, [target.x, target.y])
        
        if unit.type == '坦克':
            if dis < 2000:
                weight = 2000 - dis
                if target.type == '坦克':
                    weight += 1000
                # elif target.type == '无人机':
                    # weight += 200
                # elif target.type == '士兵':
                    # weight += 500
                else:
                    weight = -1
        elif unit.type == '步兵班':
            if dis < 1:
                weight = 800 - dis
                if target.type == '坦克':
                    weight += 500
                elif target.type == '士兵':
                    weight += 400
                else:
                    weight = -1
        elif unit.type == '无人车':
            if dis < 1000:
                weight = 1000 - dis
                if target.type == '坦克':
                    weight += 500
                elif target.type == '士兵':
                    weight += 600
                else:
                    weight = -1
        return weight
        
    def auto_attack(self, uid):
        unit = self.state.unit_dict[uid]
        if not unit.isalive:
            return 
        max_weight = -1
        tid = -1
        for target in self.state.red.qbs:
            if not target.isalive:
                continue
            weight = self.calc_weight(unit, target)
            # if unit.type == '无人车':
                # print(f"{target.type}, {target.uid}, {weight}")
            if weight > max_weight:
                tid = target.uid
                max_weight = weight
        if tid > 0:
            self.attack(uid, tid)
    
    def attack(self, uid, tid):
        unit = self.state.unit_dict[uid]
        target = self.state.unit_dict[tid]
        cmd, move_cmd = None, None
        if unit.type == '无人车':
            shell_type = -1
            cmd = FireTargetWrc(uid, tid)
            # move_cmd  = MoveToPointWrc(uid, (target.x+unit.x)/2 , (target.y+unit.y)/2, 80)
        elif unit.type == '步兵班':
            if target.type == '坦克':
                shell_type = random.choice([1,8])
            elif target.type == '无人机' : 
                shell_type = 7
            else:
                shell_type = 1
            cmd = FireTargetIcv(uid, tid, shell_type)
            # move_cmd  = MoveToPointIcv(uid, (target.x+unit.x)/2, (target.y+unit.y)/2, 80)
        elif unit.type == '坦克':
            if target.type == '坦克':
                shell_type = random.randint(1,3)
            elif target.type == '无人机':
                shell_type = 7
            else:
                shell_type = 7
            cmd = FireTargetTank(uid, tid, shell_type)
        else:
            return
        debug_print(f"{unit.type}-{uid} attack {target.type}-{tid} by shell-{shell_type}")
        debug_print(cmd)
        self.cmds.append(cmd)
        if move_cmd:
            self.cmds.append(move_cmd)

    def IcvDown(self, idx):
        if not self.isalive(IcvId[idx]):
            return
        return GetOffIcv(IcvId[idx])

class StatusTransfer:
    def __init__(self):
        print("Creat StatusTransfer Instance")
        self.trans_function_map = {
            #---------------------------------------------
            "UavInit":                  self.uav_init,
            "UavUpping":                self.uav_upping,
            "UavDetection":             self.uav_detecting,
            "UavStop":                  self.uav_stoped,
            #---------------------------------------------
            "TankInit":                 self.tank_init,
            "UacInit":                  self.uac_init,
            "IcvInit":                  self.icv_init,
            "TankGo":                   self.tank_go,
            "TankSupportWaiting":       self.tank_support_waiting,
            "TankRandomMoving":         self.tank_random_moving, 
            "UacGo":                    self.uac_go,
            "IcvGo":                    self.icv_go,
            "IcvWaiting":               self.icv_waiting,
            "IcvDown":                  self.icv_down,
            "IcvLastMoving":            self.icv_last_moving,
            #---------------------------------------------
            "AutoAttackAll":            self.auto_attack_all,
        }
        self.reset()
    
    def reset(self):
        self.state = None
        self.cross_river_check = False
        self.cross_river_check = False
        self.right_alert_check = False
        self.sold_support_check = False
        self.support_sold = False
        self.support_left = False
        self.support_right = False
    
    def transfer(self, status_list):
        if self.cross_river_check:
            self.CrossRiverCheck()
        if self.right_alert_check:
            self.RightAlertCheck()
        if self.sold_support_check:
            self.SoldSupportCheck
        new_status_list = []
        for status in status_list:
            if status.status[:3] in ['Uac', 'Uav', 'Icv', 'Tan']:
                uid = UidMap[status.status[:3]][status.idx]
                if not self.state.unit_dict[uid].isalive:
                    continue
            new_status = self.trans_function_map[status.status](status.idx)
            if new_status :
                new_status_list.append(UnitStatus(new_status, status.idx))
        return new_status_list

    #----------------
    def UavHeightOK(self, idx):
        height = self.state.unit_dict[Uavid[idx]].z
        return height > 258

    def Distance(self, uid, poi):
        unit = self.state.unit_dict[uid]
        return ((unit.x - poi[0])**2 + (unit.y - poi[1])**2)**0.5

    def UavOnPlace(self, idx):
        dis =  self.Distance(Uavid[idx], Uav_pos[idx])
        return dis < 150
        
    def UaCOnPlace(self, idx):
        dis =  self.Distance(UaCid[idx], Uac_pos[idx])
        return dis < 150

    def IcvOnPlace(self, idx):
        dis =  self.Distance(IcvId[idx], Icv_pos[idx])
        return dis < 150
        
    def TankOnPlace(self, idx):
        dis =  self.Distance(TankId[idx], Tank_pos[idx])
        return dis < 150

    def IsStoped(self, idx):
        unit = self.state.unit_dict[Uavid[idx]]
        return unit.speed < 1

    # uav-trans_function
    def uav_init(self, idx):
        if self.UavHeightOK(idx):
            return "UavDetection"
        else:
            return "UavUpping"

    def uav_upping(self, idx):
        if self.UavHeightOK(idx):
            return "UavDetection"
        else:
            return "UavUpping"

    def uav_detecting(self, idx):
        if self.UavOnPlace(idx):
            debug_print(f"无人机-{idx} is on place! Stoped!")
            return "UavStop"
        else:
            return "UavDetection"

    def uav_stoped(self, idx):
        if self.IsStoped(idx):
            return 
        return "UavStop"

    def tank_init(self, idx):
        return "TankGo"

        # tank-trans_function
    def uac_init(self, idx):
        return "UacGo"

    def icv_init(self, idx):
        return "IcvGo"
        
    def auto_attack_all(self, idx):
        return "AutoAttackAll"

    def tank_go(self, idx):
        debug_print(f"Tank-{idx} Go")
        return "TankSupportWaiting"
    
    def tank_support_waiting(self, idx):
        if self.TankOnPlace(idx):
            debug_print(f"Tank-{idx} is on place! Begin Random Move")
            return "TankRandomMoving"
        return "TankSupportWaiting"
        
    def tank_random_moving(self, idx):
        return "TankRandomMoving"
    
    def uac_go(self, idx):
        return 
    
    def icv_go(self, idx):
        return "IcvWaiting"
    
    def icv_waiting(self, idx): 
        if self.IcvOnPlace(idx):
            debug_print(f"Icv-{idx} is on place! Begin GetOff and Close Attack")
            #return "IcvDown"
            return "IcvLastMoving"
        return 'IcvWaiting'
        
    def icv_down(self, idx):
        return "IcvLastMoving"
    
    def icv_last_moving(self, idx):
        return ''

class FSMController:
    def __init__(self):
        self.action_executor = ActionGenerator()
        self.status_transfer = StatusTransfer()
        print("Creat a Controller Instance")

    def init_status(self):
        status = []
        [status.append(UnitStatus("UavInit", i)) for i in range(6)]
        [status.append(UnitStatus("TankInit", i)) for i in range(4)]
        [status.append(UnitStatus("UacInit", i)) for i in range(6)]
        [status.append(UnitStatus("IcvInit", i)) for i in range(2)]
        status.append(UnitStatus('AutoAttackAll', -1))
        self.status = status

    def reset(self):
        self.init_status()

    def step(self, state):
        self.status_transfer.state = state
        self.action_executor.state = state
        # print('------------- step ----------------')
        # for s in self.status:
            # print("Red :", s.idx, s.status)
        #print(self.status)
        self.status = self.status_transfer.transfer(self.status)
        cmds = self.action_executor.generate(self.status)
        # for cmd in cmds:
            # print(cmd)
        return cmds


def main():
    print("Exec Main Function")
    import os
    import time
    from Env.client import EnvClient
    client = EnvClient(0,0,0,0)
    client.reset()
    obs = client.get_observation()
    player = FSMController()
    player.reset()
    while True:
        cmds = player.step(obs)
        print(cmds)
        client.take_action(cmds)
        time.sleep(2)
        obs = client.get_observation()
        #obs.print_state()


if __name__ == '__main__':
    main()
